import org.apache.zookeeper.ZooKeeper;
import org.apache.zookeeper.CreateMode;
import org.apache.zookeeper.KeeperException;
import org.apache.zookeeper.KeeperException.Code;
import org.apache.zookeeper.Watcher.Event.EventType;
import org.apache.zookeeper.ZooDefs.Ids;
import org.apache.zookeeper.data.Stat;
import org.apache.zookeeper.Watcher;
import org.apache.zookeeper.WatchedEvent;

import java.util.Collections;
import java.util.HashMap;
import java.util.Iterator;
import java.util.LinkedList;
import java.util.List;
import java.util.ListIterator;
import java.util.Map;
import java.util.Queue;
import java.util.concurrent.ConcurrentLinkedQueue;

/**
 * The JobTrackerNewJobThread class
 *
 * @author  Joshua Kwan
 */

public class JobTrackerNewJobThread extends Thread {
	
	protected ZkConnector zkc;
	protected ZooKeeper zk;
	protected String userID;
	protected String hash;
	protected final int NUMPARTITIONS = FileServer.PARTITION_SIZE;
	
	protected final int SLEEPTIME = 1000;
	/**
	 * path to the job
	 */
	protected String jobPath;
	
	/**
	 * List of all the paths to the tasks in this job that have not yet been given to workers 
	 */
	protected List<String> unassignedTasks = Collections.synchronizedList(new LinkedList<String>());
//	protected Queue<String> unassignedTasks = new ConcurrentLinkedQueue<String>();
	
	//Collections.synchronizedList(new LinkedList<String>());
	
	/**
	 * Maps workers to the current tasks that have been to assigned to them for this job  
	 */
	protected Map<String, String> workerTask = Collections.synchronizedMap(new HashMap<String, String>());
	
	protected Watcher jobWatcher;
	protected Watcher wpWatcher;
	/**
	 * Maps workers to their watchers  
	 */
	protected Map<String, Watcher> workerMap = new HashMap<String, Watcher>();
	protected boolean jobFinished = false;

	/** znode paths for this class **/
	protected static final String JOBPOOL_PATH = PoolCreator.JOBPOOL_PATH;
	protected static final String WORKERPOOL_PATH = PoolCreator.WORKERPOOL_PATH;
	
	public JobTrackerNewJobThread (ZkConnector zkc, String userID, String hash){
		System.out.println("JobTrackerNewJobThread created for " + userID + " to get password for " + hash);
		this.zkc = zkc;
		this.userID = userID;
		this.hash = hash;
        zk = zkc.getZooKeeper();
        zkc.atomicAppend("Request to get password for " + hash + " accepted.", "/" + userID);
	}
	
	public void run() {
		//create a node in the /jobpool for the client
		if (zkc.exists(JOBPOOL_PATH + "/" + userID, false) == null){
        	Code dat = zkc.create(JOBPOOL_PATH + "/" + userID, null, CreateMode.PERSISTENT);
        	if (dat == Code.OK) System.out.println("JobTracker: client zNode created");
        	else System.out.println("JobTracker: client zNode already exists " + dat);
        }

		try {
			//create a new job in the /jobpool
			jobPath = zk.create(JOBPOOL_PATH+ "/" + userID + "/" + hash, null, Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
			System.out.println("zNode created: " + jobPath);
			for (int i = 0; i < NUMPARTITIONS; i++){
				String taskData = hash + ZkConnector.DATADELIMITER + Integer.toString(i);
				//System.out.println(taskData);
				synchronized(unassignedTasks){
					unassignedTasks.add(zk.create(jobPath + "/task", taskData.getBytes(), Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT_SEQUENTIAL));
				}
			}
			
			//create a watcher for the job specific to client
			jobWatcher = new Watcher(){
//				@Override
	        	public void process(WatchedEvent event) {
	        		handleJobEvent(event);
	        	}
			};
						
			// attaching job watcher to client job
			checkPath(jobPath, jobWatcher);

			//create the watcher for the worker pool
			wpWatcher = new Watcher(){
//	        	@Override
	        	public void process(WatchedEvent event) {
	        		handleWorkerPoolEvent(event);
	        	}
	        };
	        //get the children of the workerpool
	        System.out.println("getting the children of the worker pool");
	        List<String> wpList  = zkc.getZooKeeper().getChildren(PoolCreator.WORKERPOOL_PATH, wpWatcher);
	        System.out.println(WORKERPOOL_PATH + " contains " + wpList.size());
	        updateWorkerMap(wpList);
	        
	        initialAssignment();

	        while (!jobFinished) {
	        	Thread.sleep(SLEEPTIME);
	        	for(String wp: workerMap.keySet()){
	        		checkWorker(wp);
	        	}
	        }
	        
		} catch (KeeperException e) {
			System.out.println(e);
			if (e.code() == KeeperException.Code.NODEEXISTS){
				System.out.println("A job for " + hash + " has already been requested.");
				//exit?
			}
		} catch (InterruptedException e) {
			System.out.println(e);
			//try again?...
		} catch (IllegalArgumentException e) {
			System.out.println(e);
			//maybe exit??
		}
	}
	

	protected Stat checkPath(String path, Watcher watcher) {
        Stat stat = zkc.exists(path, watcher);
        return stat;
    }

    protected void handleJobEvent(WatchedEvent event){
        String path = event.getPath();
        EventType type = event.getType();
        if(path.equalsIgnoreCase(jobPath)) {
            if (type == EventType.NodeDataChanged) {
            	//if there is a result, empty the unassigned list
            	try {
            		Stat stat = new Stat();
            		byte[] jobData = zk.getData(path, jobWatcher, stat); 
            		if (jobData != null){ 
                		unassignedTasks.clear();
                	}
            	} catch (InterruptedException e) {
            		System.out.println(e);
            	} catch (KeeperException e) {
            		System.out.println(e);
            	}
            	
            } else {
            	checkPath(path, jobWatcher);
            }
        }
    }
	
    protected void handleWorkerPoolEvent(WatchedEvent event) {
        String path = event.getPath();
        EventType type = event.getType();
        if(path.equalsIgnoreCase(WORKERPOOL_PATH)) {
            if (type == EventType.NodeChildrenChanged) {
            	try {
            		List<String> wpList = zk.getChildren(WORKERPOOL_PATH, wpWatcher);
                	updateWorkerMap(wpList);	
            	} catch (InterruptedException e) {
            		System.out.println(e);
            	} catch (KeeperException e) {
            		System.out.println(e);
            	}
            } else {
            	checkPath(path, wpWatcher);
            }
        }
    }
    
    protected void handleWorkerEvent(WatchedEvent event) {
    	String path = event.getPath();
        EventType type = event.getType();
        if (workerMap.containsKey(path)){
    	    if (type == EventType.NodeDataChanged){
   	    		checkWorker(path);
    	    } else if (type == EventType.NodeDeleted){
    	    	//see if the worker was processing a task for us
    	    	//if so, move that task from workerTask back to unassignedTasks
    	    	if (workerTask.containsKey(path)) {
    	    		String taskPath = workerTask.get(path);
    	    		if (zkc.exists(taskPath, false) != null) {
    	    			unassignedTasks.add(taskPath);
    	    			workerTask.remove(path);
    	    			//remove the path from the data of taskPath
    	    			atomicRemoveTaskPath(path, taskPath);
    	    		}
    	    	}
    	    } else {
    	    	checkPath(path, workerMap.get(path));
    	    }
        }
    }
    
    /**
     * Check if to see that the workermap is still up to date 
     * (i.e. handles case where Worker crashes or new workers 
     * are added),distribution of unassigned(crashed) tasks.
     * 
     * @param workerList
     */
    protected void updateWorkerMap (List<String> workerList) {

    	//check paths in workerList & create if not found in workerMap  
    	// Existing workers against what you are holding now
    	for(String path: workerList){
    		if(!workerMap.containsKey(path)){
    			System.out.println(PoolCreator.WORKERPOOL_PATH +  "/" + path);
    			newWorker(PoolCreator.WORKERPOOL_PATH + "/" + path);
    			//TODO: add the task assigned to this worker to unassigned task
    		}
    	}

    	// removal of dead workers from our workerMap 
    	// check entries in workerMap & remove if not found
    	for(String wme: workerMap.keySet()){
			System.out.println(wme.substring(WORKERPOOL_PATH.lastIndexOf("/") + 1) + " not found!");
			if (checkPath(wme, workerMap.get(wme)) == null){
				workerMap.remove(wme);
			}
			//TODO: add the task assigned to this worker to unassigned task
    	}
    	System.out.println("WorkerMap size=" + workerMap.size() + " unassigned tasks=" + unassignedTasks.size());
    }
    
    /**
     * checks the worker data and attempts to assign a task to the worker if the data is empty
     * @param path - path to the worker
     * @throws KeeperException
     * @throws InterruptedException
     */
    protected void checkWorker(String path) {
    	// get data and respond to new requests
    	System.out.println("checking worker at " + path + " unassigned tasks=" + unassignedTasks.size() );
 		Stat stat = new Stat();
 		try {
 			byte [] workerData = zk.getData(path, workerMap.get(path), stat);
 	 		if (workerMap.get(path) == null) System.out.println("Watcher for " + path + " is null!!");
 			if (workerData == null || workerData.length == 0) {
 				//try to give it a new request
 				synchronized(unassignedTasks){
 					if (!unassignedTasks.isEmpty()){
 						String taskPath = unassignedTasks.get(0);
 						System.out.println("trying to assign " + taskPath + " to " + path);
 	    				if (assignTask(path, taskPath, stat.getVersion())){
 	    					taskPath = unassignedTasks.remove(0);
 	    					checkPath(path, workerMap.get(path));
 	    					synchronized(workerTask){
 	    						workerTask.put(path, taskPath);
 	    					}
 	    				} else {
 	    					//check to see if the worker was previously working on one of our tasks. If so, remove this entry from workerTask
 	    					if (workerTask.containsKey(path)) {
 	    						workerTask.remove(path);
 	    					}
 	    				}
 					} else {
 						//check the job for results and check tasks in this job for completion
 						List<String> taskList = zk.getChildren(jobPath, false);
 						if (taskList.size() == 0) {
 							//job completed, delete all watches on the workerPool
 							workerMap.clear();
 							jobFinished = true;
 						}
 					}
 				}
 			} /*else {
 				System.out.println(path + " is busy, worker data is " + new String(workerData));
 			}*/
 		} catch (InterruptedException e) {
 			System.out.println(e);
 		} catch (KeeperException e) {
 			System.out.println(e);
 		}
    }
    
    /**
     * attempts to assign a task to the idle worker
     * @param workerPath - path to the worker
     * @param taskPath - path to the task
     * @param version - version of most recently read data
     * @return true if successfully assigned task, false if assign failed
     */
    protected boolean assignTask(String workerPath, String taskPath, int version) {
    	try {
    		System.out.println("attempting to assign task to worker");
    		zk.setData(workerPath, taskPath.getBytes(), version);
    		zkc.atomicAppend(workerPath, taskPath);
    	} catch (KeeperException e) {
    		return false;
    	} catch (InterruptedException e) {
    		return false;
    	}
    	return true;
    }
    
    protected void initialAssignment(){
    	for (Iterator wmi = workerMap.keySet().iterator(); wmi.hasNext() && !unassignedTasks.isEmpty(); ) {
    		String workerPath = (String) wmi.next();
    		System.out.println("Info: " + workerPath);
    		checkWorker(workerPath);
    	}
    }
    
    protected void newWorker(String path) {
    	Watcher wWatcher = new Watcher(){
//    		@Override
        	public void process(WatchedEvent event) {
        		handleWorkerEvent(event);
        	}
    	};
    	workerMap.put(path, wWatcher);
    	checkWorker(path);
    }
    
    /**
     * removes the workerPath from the the data in path
     * @param workerPath - path to the worker
     * @param taskPath - path to the task
     * @return stat of the zNode
     */
    protected Stat atomicRemoveTaskPath(String workerPath, String taskPath){
    	while (true){
        	try{
        		Stat stat = new Stat();
        		byte [] currentData = zk.getData(taskPath, false, stat);
        		if(currentData != null){
        			String strCurrentData = new String(currentData);
        			String[] data = strCurrentData.split(ZkConnector.DATADELIMITER); 
        			if (data.length == 3 && data[2].equals(workerPath)){
        				zk.setData(taskPath, strCurrentData.substring(0, strCurrentData.lastIndexOf(ZkConnector.DATADELIMITER)).getBytes(), stat.getVersion());
        			}
        		} 
        		return stat;
        	} catch (KeeperException e){
        		if (e.code() == KeeperException.Code.NONODE) return null;
        		//we could also get KeeperException.BadVersion, but we just try again if unsuccessful
        	} catch (InterruptedException e){
        		//well we just try again if unsuccessful anyways...
        	}
    	}
    }
}